package com.zhangx.datastructure.search;

import com.zhangx.datastructure.search.dto.BiNode;

/**
 * @ClassName AVLTree
 * @Description 平衡二叉树
 * @Author zhangx
 * @Date 2021/1/29 8:55
 **/
public class AVLTree {

    private static final int LH = 1;//左高
    private static final int EH = 0;//等高
    private static final int RH = -1;//右高

    private BiNode tree = null;

    /**
     * 以p为根
     * 向右旋转
     *
     * @param p
     */
    void rightRotate(BiNode p) {
        BiNode l = p.getLchild();//l 指向P的左子树根节点
        p.setLchild(l.getRchild());//l右子树挂到P的左子树
        l.setRchild(p);//
        p = l; //p指向新的根节点
    }

    /**
     * 以p为根
     * 向左旋转
     *
     * @param p
     */
    void leftRotate(BiNode p) {
        BiNode r = p.getRchild();//r 指向P的右子树根节点
        p.setRchild(r.getLchild());//r左子树挂到P的右子树
        r.setLchild(p);//
        p = r; //p指向新的根节点
    }

    /**
     * 对t 为根的二叉树进行左平衡旋转处理
     * 左子树高于右子树，BF>1
     * 结束时t 指向新的根节点
     *
     * @param t
     */
    void leftBalance(BiNode t) {
        BiNode l, lr;
        l = t.getLchild();//l指向t 的左子树
        switch (l.getBf()) {//检查T的平衡度进行处理
            case LH://新节点插入在t的左孩子的左子树，要做单右旋处理
                t.setBf(EH);
                l.setBf(EH);
                rightRotate(t);
                break;
            case RH://新节点插入在t的左孩子的右子树，要做双旋处理
                lr = l.getRchild();
                switch (lr.getBf()) {//修改t 及其左孩子的平衡因子
                    case LH:
                        t.setBf(RH);
                        l.setBf(EH);
                        break;
                    case EH:
                        t.setBf(EH);
                        l.setBf(EH);
                        break;
                    case RH:
                        t.setBf(EH);
                        l.setBf(LH);
                        break;
                }
                lr.setBf(EH);
                leftRotate(t.getLchild());//t 的左子树做左平衡
                rightRotate(t);//t 做右旋平衡处理
        }
    }


    /*  对以指针T所指结点为根的二叉树作右平衡旋转处理， */
    /*  本算法结束时，指针T指向新的根结点 */
    void rightBalance(BiNode t) {
        BiNode r, rl;
        r = t.getRchild(); /*  R指向T的右子树根结点 */
        switch (r.getBf()) { /*  检查T的右子树的平衡度，并作相应平衡处理 */
            case RH: /*  新结点插入在T的右孩子的右子树上，要作单左旋处理 */
                t.setBf(EH);
                r.setBf(EH);
                leftRotate(t);
                break;
            case LH: /*  新结点插入在T的右孩子的左子树上，要作双旋处理 */
                rl = r.getLchild();/*  Rl指向T的右孩子的左子树根 */
                switch (rl.getBf()) { /*  修改T及其右孩子的平衡因子 */
                    case RH:
                        t.setBf(LH);
                        r.setBf(EH);
                        break;
                    case EH:
                        t.setBf(EH);
                        r.setBf(EH);
                        break;
                    case LH:
                        t.setBf(EH);
                        r.setBf(RH);
                        break;
                }
                rl.setBf(EH);
                rightRotate(t.getRchild()); /*  对T的右子树作右旋平衡处理 */
                leftRotate(t); /*  对T作左旋平衡处理 */
        }
    }
    /*  若在平衡的二叉排序树T中不存在和e有相同关键字的结点，则插入一个 */
    /*  数据元素为e的新结点，并返回1，否则返回0。若因插入而使二叉排序树 */
    /*  失去平衡，则作平衡旋转处理，布尔变量taller反映T长高与否。 */
    private Boolean insertAVL(BiNode t, int e, Boolean taller) {
        if (t==null){
            /*  插入新结点，树“长高”，置taller为TRUE */
            t = new BiNode();
            t.setData(e);
            t.setBf(EH);
            taller = true;
        }else {
            if (e == t.getData()){ /*  树中已存在和e有相同关键字的结点则不再插入 */
                taller = false;
                return false;
            }
            if (e<t.getData()){//当e 小于t 根节点，要到他的左子树查找
                if (!insertAVL(t.getLchild(),e,taller)){//未插入
                    return false;
                }
//                插入成功
                if (taller){/*   已插入到T的左子树中且左子树“长高” */
                    switch (t.getBf()){
                        case LH: /*  原本左子树比右子树高，需要作左平衡处理 */
                            leftBalance(t);
                            taller=false;
                            break;
                        case EH: /*  原本左、右子树等高，现因左子树增高而使树增高 */
                            t.setBf(LH);
                            taller=true;
                            break;
                        case RH: /*  原本右子树比左子树高，现左、右子树等高 */
                            t.setBf(EH);
                            taller=false;
                            break;

                    }
                }
            }else {
                /*  应继续在T的右子树中进行搜索 */
                if(!insertAVL(t.getRchild(),e,taller)) { /*  未插入 */
                    return false;
                }
                if (taller){/*  已插入到T的右子树且右子树“长高” */
                    switch(t.getBf()) /*  检查T的平衡度 */
                    {
                        case LH: /*  原本左子树比右子树高，现左、右子树等高 */
                            t.setBf(EH);
                            taller=false;
                            break;
                        case EH: /*  原本左、右子树等高，现因右子树增高而使树增高  */
                            t.setBf(RH);
                            taller=true;
                            break;
                        case RH: /*  原本右子树比左子树高，需要作右平衡处理 */
                            rightBalance(t);
                            taller=false;
                            break;
                    }

                }

            }
        }


        return true;
    }

    public static void main(String[] args) {
        AVLTree avlTree = new AVLTree();
        int[] a= new int[]{3,2,1,4,5,6,7,10,9,8};
//        BiNode t = null;
        Boolean taller = false;
        for(int i=0;i<10;i++)
        {
            avlTree.insertAVL(avlTree.tree,a[i],taller);
        }
        System.out.println(avlTree.tree);

    }


}
